"""
日志管理模块

提供统一的日志记录功能，支持控制台输出和文件保存
"""

import logging
import sys
from pathlib import Path
from typing import Optional, Union
from datetime import datetime
import torch


class Logger:
    """统一的日志管理器"""
    
    def __init__(self, 
                 name: str = "SAM-Echo-Video",
                 log_dir: Optional[Union[str, Path]] = None,
                 level: int = logging.INFO,
                 console_output: bool = True,
                 file_output: bool = True):
        """
        初始化日志器
        
        Args:
            name: 日志器名称
            log_dir: 日志文件保存目录
            level: 日志级别
            console_output: 是否输出到控制台
            file_output: 是否保存到文件
        """
        self.logger = logging.getLogger(name)
        self.logger.setLevel(level)
        
        # 清除已有的处理器
        self.logger.handlers.clear()
        
        # 设置格式
        formatter = logging.Formatter(
            '%(asctime)s - %(name)s - %(levelname)s - %(message)s',
            datefmt='%Y-%m-%d %H:%M:%S'
        )
        
        # 控制台输出
        if console_output:
            console_handler = logging.StreamHandler(sys.stdout)
            console_handler.setLevel(level)
            console_handler.setFormatter(formatter)
            self.logger.addHandler(console_handler)
        
        # 文件输出
        if file_output and log_dir is not None:
            log_dir = Path(log_dir)
            log_dir.mkdir(parents=True, exist_ok=True)
            
            # 创建带时间戳的日志文件名
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            log_file = log_dir / f"{name}_{timestamp}.log"
            
            file_handler = logging.FileHandler(log_file, encoding='utf-8')
            file_handler.setLevel(level)
            file_handler.setFormatter(formatter)
            self.logger.addHandler(file_handler)
    
    def info(self, message: str) -> None:
        """记录信息日志"""
        self.logger.info(message)
    
    def debug(self, message: str) -> None:
        """记录调试日志"""
        self.logger.debug(message)
    
    def warning(self, message: str) -> None:
        """记录警告日志"""
        self.logger.warning(message)
    
    def error(self, message: str) -> None:
        """记录错误日志"""
        self.logger.error(message)
    
    def critical(self, message: str) -> None:
        """记录严重错误日志"""
        self.logger.critical(message)
    
    def log_experiment(self, 
                      experiment_name: str,
                      config: dict,
                      metrics: dict,
                      epoch: Optional[int] = None) -> None:
        """记录实验信息"""
        self.info(f"实验: {experiment_name}")
        if epoch is not None:
            self.info(f"Epoch: {epoch}")
        self.info(f"配置: {config}")
        self.info(f"指标: {metrics}")
    
    def log_training(self, 
                    epoch: int,
                    loss: float,
                    metrics: dict,
                    lr: Optional[float] = None) -> None:
        """记录训练信息"""
        message = f"Epoch {epoch:3d} - Loss: {loss:.6f}"
        if lr is not None:
            message += f" - LR: {lr:.2e}"
        
        for key, value in metrics.items():
            if isinstance(value, (int, float)):
                message += f" - {key}: {value:.4f}"
            else:
                message += f" - {key}: {value}"
        
        self.info(message)
    
    def log_evaluation(self, 
                      dataset_name: str,
                      metrics: dict) -> None:
        """记录评估信息"""
        self.info(f"评估数据集: {dataset_name}")
        for key, value in metrics.items():
            if isinstance(value, (int, float)):
                self.info(f"  {key}: {value:.4f}")
            else:
                self.info(f"  {key}: {value}")
    
    def log_model_info(self, model: torch.nn.Module) -> None:
        """记录模型信息"""
        total_params = sum(p.numel() for p in model.parameters())
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        
        self.info(f"模型总参数量: {total_params:,}")
        self.info(f"可训练参数量: {trainable_params:,}")
        self.info(f"模型结构: {model.__class__.__name__}")
    
    def log_system_info(self) -> None:
        """记录系统信息"""
        self.info(f"PyTorch版本: {torch.__version__}")
        self.info(f"CUDA可用: {torch.cuda.is_available()}")
        if torch.cuda.is_available():
            self.info(f"CUDA版本: {torch.version.cuda}")
            self.info(f"GPU数量: {torch.cuda.device_count()}")
            for i in range(torch.cuda.device_count()):
                self.info(f"GPU {i}: {torch.cuda.get_device_name(i)}")


# 全局日志器实例
logger = Logger()


def get_logger(name: str = "SAM-Echo-Video") -> Logger:
    """获取日志器实例"""
    return Logger(name)


def setup_logging(log_dir: Optional[Union[str, Path]] = None,
                  level: int = logging.INFO) -> Logger:
    """设置项目日志"""
    return Logger(log_dir=log_dir, level=level)

